import matplotlib as mpl
mpl.rcParams['font.sans-serif'] = ['SimHei']
mpl.rcParams['axes.unicode_minus'] = False
from sklearn.model_selection import train_test_split

def load_data():
    from sklearn.datasets import make_classification

    X, y = make_classification(
        n_samples=1000,
        n_features=10,
        n_informative=3,
        n_redundant=0,
        n_repeated=0,
        n_classes=2,
        random_state=0,
        shuffle=False,
    )
    return  X, y



if __name__ == '__main__':
    X, y = load_data()
    # 拆分数据
    X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y, random_state=42)
    # 随机森林训练
    from sklearn.ensemble import RandomForestClassifier
    feature_names = [f"feature {i}" for i in range(X.shape[1])]
    forest = RandomForestClassifier(random_state=0)
    forest.fit(X_train, y_train)

    # 1. 特征重要性 impurity-based importance
    import time
    import numpy as np
    start_time = time.time()
    importances = forest.feature_importances_
    std = np.std([tree.feature_importances_ for tree in forest.estimators_], axis=0)
    elapsed_time = time.time() - start_time

    print(f"计算重要性所用的时间: {elapsed_time:.3f} seconds")

    # 特征重要性可视化
    import pandas as pd
    import matplotlib.pyplot as plt
    forest_importances = pd.Series(importances, index=feature_names)
    fig = plt.figure()
    ax = fig.add_subplot(121)
    forest_importances.plot.bar(yerr=std, ax=ax)
    ax.set_title("使用MDE计算特征重要性")
    ax.set_ylabel("不纯度的平均减少量")
    fig.tight_layout()

    # 2. Permutation feature importance 它们不偏向于高基数特征，可以在遗漏的测试集上计算
    from sklearn.inspection import permutation_importance
    start_time = time.time()
    result = permutation_importance(
        forest, X_test, y_test, n_repeats=10, random_state=42, n_jobs=2
    )
    elapsed_time = time.time() - start_time
    print(f"计算重要性所用的时间: {elapsed_time:.3f} seconds")
    forest_importances = pd.Series(result.importances_mean, index=feature_names)
    ax = fig.add_subplot(122)
    forest_importances.plot.bar(yerr=result.importances_std, ax=ax)
    ax.set_title("使用排列计算特征重要性")
    ax.set_ylabel("平均精度下降")
    fig.tight_layout()
    plt.show()

